// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/experimental/wave/dwio/ColumnReader.h>
#include <pollux/experimental/wave/exec/WaveSplitReader.h>
#include <pollux/experimental/wave/exec/tests/utils/FileFormat.h>

namespace kumo::pollux::wave::test {

/// A WaveSplitReader that decodes mock Wave tables.
class WaveTestSplitReader : public WaveSplitReader {
 public:
  WaveTestSplitReader(
      const std::shared_ptr<connector::ConnectorSplit>& split,
      const SplitReaderParams& params,
      const DefinesMap* defines);

  bool emptySplit() override {
    return !stripe_ || stripe_->columns[0]->numValues == 0 || emptySplit_;
  }

  int32_t canAdvance(WaveStream& stream) override;

  void schedule(WaveStream& stream, int32_t maxRows = 0) override;

  vector_size_t outputSize(WaveStream& stream) const;

  bool isFinished() const override;

  uint64_t getCompletedBytes() override {
    return params_.ioStats->rawBytesRead();
  }

  uint64_t getCompletedRows() override {
    return nextRow_;
  }

    std::unordered_map<std::string, RuntimeCounter> runtimeStats() override {
    return {};
  }

  static void registerTestSplitReader();

 private:
  int32_t available() const {
    return stripe_->columns[0]->numValues - nextRow_;
  }

  std::shared_ptr<connector::ConnectorSplit> split_;
  SplitReaderParams params_;
  FileHandleCachedPtr fileHandleCachePtr;
  cache::AsyncDataCache* cache_{nullptr};
  test::Stripe* stripe_{nullptr};
  std::unique_ptr<ColumnReader> columnReader_;
  // First unscheduled row.
  int32_t nextRow_{0};
  int32_t scheduledRows_{0};
  dwio::common::ColumnReaderStatistics readerStats_;
  raw_vector<int32_t> rows_;
  FileHandleCachedPtr fileHandle_;
  FileInfo fileInfo_;
  bool emptySplit_{false};
};

} // namespace kumo::pollux::wave::test
